from time import sleep
import os
import sys
import matplotlib.pyplot as plt
import torch
import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader
from tqdm import tqdm
import math
import time
import numpy as np
from torch.utils.tensorboard import SummaryWriter

class Valer:
    def __init__(self, config, loader):
        self.writer = SummaryWriter(f"./project_datasave/{config['project_name']}/log")
        self.config = config
        self.train_cfg = config["trainer"]
        if (self.train_cfg["num_sub_heads"] == None):
            self.num_sub_heads = 1
        else:
            self.num_sub_heads = self.train_cfg["num_sub_heads"]

        self.val_time = 0
        self.attack_time = 0

        # 用于计算输出分类类别顺序
        self.pre_num = [[0 for i in range(self.config["datasets"]["label_type"])] for j in range(self.num_sub_heads)]
        self.label_num = [0 for i in range(self.config["datasets"]["label_type"])]

        self.key_score = dict()
        """ 密钥—分数字典 """
        self.mval_key = dict()
        """ 中间值-密钥字典 """

        self.order = [[j for j in range(self.config["datasets"]["label_type"])] for i in range(self.num_sub_heads)]
        self.val_loader = loader
        self.device = self.config["trainer"]["device"]
        self.set_net()

    def set_net(self):
        net_structure = self.train_cfg["net"]
        print("将使用%s网络结构进行验证" % net_structure)
        if (net_structure == 'cs3'):
            from Nets.cnn_single_head_3layer import CNNNet
            self.model = CNNNet(self.config).to(self.device)
        elif (net_structure == 'resnet18'):
            from Nets.Resnet import ResNet_18
            self.model = ResNet_18(self.config).to(self.device)
        elif (net_structure == 'cs4'):
            from Nets.cnn_single_head import CNNNet
            self.model = CNNNet(self.config).to(self.device)
        elif (net_structure == 'ms5'):
            from Nets.mlp_5layer import MLPNet
            self.model = MLPNet(self.config).to(self.device)
        elif (net_structure == 'cm3'):
            from Nets.cnn_multi_head import CNNNet
            self.model = CNNNet(self.config).to(self.device)
        elif (net_structure == 'mm5'):
            from Nets.mlp_multi_head_5layer import MLPNet
            self.model = MLPNet(self.config).to(self.device)
        elif (net_structure == 'mm7'):
            from Nets.mlp_multi_head_7layer import MLPNet
            self.model = MLPNet(self.config).to(self.device)
        elif (net_structure == 'cmp3'):
            from Nets.cnn_multi_head_pca import CNNNet
            self.model = CNNNet(self.config).to(self.device)
        elif (net_structure == 'csp3'):
            from Nets.cnn_single_head_pca import CNNNet
            self.model = CNNNet(self.config).to(self.device)

    def load_weight(self, model_path):
        self.weights = model_path.split("/")[-1][:-4]
        self.model.load_state_dict(torch.load(model_path))
        self.model.eval()

    # 计算单个模型准确率
    def get_acc(self):
        print("计算准确率......")

        correct = [0 for i in range(self.num_sub_heads)]
        accuracy = [0 for i in range(self.num_sub_heads)]
        total = 0
        self.model.eval()
        with torch.no_grad():
            for data in tqdm(self.val_loader):
                input, _, labels = data
                input, labels = input.to(self.device), labels.to(self.device)
                for l in labels:
                    self.label_num[l] += 1
                outputs = self.model(input)
                total += labels.size(0)
                if (isinstance(outputs, list)):
                    for i in range(self.train_cfg["num_sub_heads"]):
                        _, predicted = torch.max(outputs[i].data, 1)
                        for ip, p in enumerate(predicted):
                            predicted[ip] = self.order[i][p]
                            self.pre_num[i][p] += 1
                        correct[i] += (predicted == labels).sum().item()
                else:
                    _, predicted = torch.max(outputs.data, 1)
                    for ip, p in enumerate(predicted):
                        predicted[ip] = self.order[0][p]
                        self.pre_num[0][p] += 1
                    correct[0] += (predicted == labels).sum().item()
        for i in range(self.num_sub_heads):
            accuracy[i] = correct[i] / total
            print(f"{i+1}. 正确个数:{correct[i]}/{total} 测试总准确率: {accuracy[i] * 100:.2f}%")
        return accuracy

    # 计算单个模型准确率
    def get_matrix(self):
        matrix = [[[0 for i in range(self.config["datasets"]["label_type"])] for j in range(self.config["datasets"]["label_type"])] for k in range(self.num_sub_heads)]
        self.model.eval()
        with torch.no_grad():
            for data in tqdm(self.val_loader):
                input, _, labels = data
                input, labels = input.to(self.device), labels.to(self.device)
                outputs = self.model(input)
                if (isinstance(outputs, list)):
                    for i in range(self.train_cfg["num_sub_heads"]):
                        _, predicted = torch.max(outputs[i].data, 1)
                        for ip, p in enumerate(predicted):
                            predicted[ip] = self.order[i][p]
                        for l, p in zip(labels, predicted):
                            matrix[i][p][l] += 1
                else:
                    _, predicted = torch.max(outputs.data, 1)
                    for ip, p in enumerate(predicted):
                        predicted[ip] = self.order[0][p]
                    for l, p in zip(labels, predicted):
                        matrix[0][p][l] += 1

        # 写入文件 #
        for i in range(self.num_sub_heads):
            fig, ax = plt.subplots()
            with open(f"./project_datasave/{self.config['project_name']}/log/matrix_{self.weights}_head{i}.txt", "w") as f:
                for m in matrix[i]:
                    for n in m:
                        f.write(f"{n} ")
                    f.write("\n")
            cax = ax.matshow(matrix[i], cmap='viridis')
            for k in range(np.array(matrix[i]).shape[0]):
                for j in range(np.array(matrix[i]).shape[1]):
                    ax.text(j, k, str(matrix[i][k][j]), ha='center', va='center', color='white')
            fig.colorbar(cax)
            fig.savefig(f"./project_datasave/{self.config['project_name']}/figure/matrix_{self.weights}_head{i}.png")
            self.writer.add_figure(f"matrix_{self.weights}_head{i}", fig)


    def get_model_num(self):
        print("获取最后权重分类顺序......")
        m = "./project_datasave/" + self.config["project_name"] + "/model/every_epochs/model_" + str(self.config["trainer"]["epochs"]-1) + ".pth"
        self.load_weight(m)
        self.get_acc()

    # 获取分类顺序
    def get_class_order(self):
        for i in range(self.num_sub_heads):
            label_num = self.label_num.copy()
            for j in range(self.config["datasets"]["label_type"]):
                self.order[i][self.pre_num[i].index(max(self.pre_num[i]))] = label_num.index(max(label_num))
                self.pre_num[i][self.pre_num[i].index(max(self.pre_num[i]))] = -1
                label_num[label_num.index(max(label_num))] = -1
        
    # 计算所有模型最高准确率
    def highest_acc(self):
        acc_all = []
        highest_acc = 0.0
        highest_model = ""
        highest_head = 0
        m = os.listdir("./project_datasave/" + self.config["project_name"] + "/model/every_epochs")
        m.sort(key=lambda x: int(x.split("_")[1].split(".")[0]))
        m_path = "./project_datasave/" + self.config["project_name"] + "/model/every_epochs/"
        for path in m:
            self.load_weight(m_path + path)
            acc = self.get_acc()
            acc_all.append(acc)
            for e, i in enumerate(acc):
                with open("./project_datasave/" + self.config["project_name"] + f"/log/acc_head{e}.txt", "a") as f:
                    f.write(f"{i}\n")
                if (i > highest_acc):
                    highest_acc = i
                    highest_model = path
                    highest_head = e
        print(f"最高准确率模型: {highest_model} 准确率: {highest_acc * 100:.2f}% 第{highest_head+1}个子头")
        # 绘制图像 #
        fig, ax = plt.subplots()
        for i in range(self.num_sub_heads):
            ax.plot([j[i] for j in acc_all])
        ax.set_xlabel("epochs")
        ax.set_ylabel("accuracy")
        fig.savefig(f"./project_datasave/{self.config['project_name']}/figure/acc.png")
        self.writer.add_figure("acc", fig)
        self.hightest_model = highest_model
        self.hightest_head = highest_head

    def cal_middleval(self, p_text, key):
        def hamming_weight(num):
            count = 0
            while num:
                count += 1
                num &= num - 1
            return count
        """ 计算sbox(p[3] xor k[3])
        注意要传两个数字型
        """
        s_box = [[0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76],
                 [0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0],
                 [0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15],
                 [0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75],
                 [0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84],
                 [0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF],
                 [0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8],
                 [0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2],
                 [0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73],
                 [0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB],
                 [0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79],
                 [0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08],
                 [0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A],
                 [0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E],
                 [0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF],
                 [0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16]]
    
        m_val = p_text ^ key
        row = m_val >> 4
        col = m_val & 0x0F
        if (self.config["datasets"]["label_type"] == 2):
            mid_result = int(s_box[row][col]) >> 7
        elif (self.config["datasets"]["label_type"] == 9):
            mid_result = hamming_weight(int(s_box[row][col]))
        elif (self.config["datasets"]["label_type"] == 256):
            mid_result = int(s_box[row][col])
        return mid_result

    def attack_val(self):
        print("攻击验证......")
        self.model.eval()
        for n in range(self.num_sub_heads):
            rank_list = []
            for i in range(256):
                self.key_score[i] = 0.0
            print(f"head{n}...")
            with torch.no_grad():
                for i, data in enumerate(tqdm(self.val_loader)):
                    att_trace, att_plt, att_label = data
                    att_trace = att_trace.to(self.device)
                    att_plt = att_plt.to(self.device)
                    att_label = att_label.to(self.device)

                    outputs = self.model(att_trace)

                    # 中间值分数 #
                    if (isinstance(outputs, list)):
                        softmax_outputs_batch = outputs[n]
                    else:
                        softmax_outputs_batch = outputs

                    for softmax_output, plain_texts in zip(softmax_outputs_batch, att_plt):
                        # 生成中间值-密钥字典 #
                        # 二分类
                        for o in range(self.config["datasets"]["label_type"]):
                            self.mval_key[o] = []
                        for k in range(256):
                            self.mval_key[self.cal_middleval(plain_texts, k)].append(k)

                        # 将中间值分数赋给密钥 #
                        for mv in range(self.config["datasets"]["label_type"]):
                            # 按照output_order的顺序换掉softmax_output的顺序
                            output_temp = softmax_output.clone()
                            for j in range(len(self.order[n])):
                                softmax_output[self.order[n][j]] = output_temp[j]
                            for k in self.mval_key[mv]:
                                a = self.key_score[k]
                                b = softmax_output[mv].item()
                                if (b != 0):
                                    self.key_score[k] = a + math.log(b, 2)
                                else:
                                    self.key_score[k] = a - 9999999

                        # 对字典进行排序 #
                        key_score_ = sorted(self.key_score.items(), key=lambda d: d[1], reverse=True)

                        # 查找224的排名
                        for j in range(256):
                            if key_score_[j][0] == 224:
                                rank = j
                                break
                        rank_list.append(rank)

            # 写入文件 #
            with open(f"./project_datasave/{self.config['project_name']}/log/rank_{self.weights}_head{n}.txt", "w") as f:
                for rank in rank_list:
                    f.write(f"{rank}\n")

            # 绘制图像 #
            fig, ax = plt.subplots()
            ax.plot(rank_list)
            ax.set_xlabel("traces")
            ax.set_ylabel("rank")
            fig.savefig(f"./project_datasave/{self.config['project_name']}/figure/rank_{self.weights}_head{n}.png")
            self.writer.add_figure(f"rank_{self.weights}_head{n}", fig)


    def val(self):
        if (self.train_cfg["type"] != "supervised"):
            self.get_model_num()
            print(f"pre_num:{self.pre_num}")
            print(f"label_num:{self.label_num}")
            self.get_class_order()
            print(f"order: {self.order}")
        print("判定最好的模型......")
        val_start_time = time.time()
        self.highest_acc()
        val_end_time = time.time()
        self.val_time = val_end_time - val_start_time
        self.load_weight("./project_datasave/" + self.config["project_name"] + "/model/every_epochs/" + self.hightest_model)
        print("加载最好的模型......")
        attact_start_time = time.time()
        self.attack_val()
        attack_end_time = time.time()
        self.attack_time = attack_end_time - attact_start_time
        self.get_matrix()
        self.load_weight("./project_datasave/" + self.config["project_name"] + "/model/every_epochs/" + "model_" + str(self.config["trainer"]["epochs"]-1) + ".pth")
        print("加载最后的模型......")
        self.attack_val()
        self.get_matrix()
